from bpfcc import BPF, USDT
import json
import time
import datetime
import sys

# 定义BPF程序
bpf_text = """
#include <uapi/linux/sched.h>
#include <linux/ptrace.h>

struct offcpu_event_t {
    u32 pid;
    u64 start_time;
    u64 end_time;
};

BPF_HASH(offcpu_maps, u32, struct offcpu_event_t);
BPF_PERF_OUTPUT(events);


static int start(struct pt_regs *ctx, struct task_struct *prev) {
    u32 prev_pid = prev->pid;
    struct offcpu_event_t new_evt = {0};

    // 检查是否为目标PID
    if (prev_pid != TARGET_PID) {
        return 0;
    }
	
    u64 ts = bpf_ktime_get_ns();

    new_evt.start_time = ts;
    new_evt.pid = prev_pid;
    offcpu_maps.update(&prev_pid, &new_evt);
    return 0;
}

static int stop(struct pt_regs *ctx, struct task_struct *prev) {
    u32 cur_pid = (u32)bpf_get_current_pid_tgid();
    struct offcpu_event_t *evt;
	
    u64 ts = bpf_ktime_get_ns();

    evt = offcpu_maps.lookup(&cur_pid);
    if (evt) {
	evt->end_time = ts;
	events.perf_submit(ctx, evt, sizeof(struct offcpu_event_t));
	offcpu_maps.delete(&cur_pid);
    }
    return 0;
}

int trace_sched_switch(struct pt_regs *ctx, struct task_struct *prev) {
    start(ctx, prev);
    stop(ctx, prev);
    return 0;
}

"""

def _convert_ktime_to_unix(bpf_ktime_ns):
    # 获取当前时间戳（CLOCK_REALTIME）和系统运行时间（uptime）
    current_realtime = time.time()
    with open("/proc/uptime", "r") as f:
        uptime_seconds = float(f.read().split()[0])
    
    # 计算系统启动时间（CLOCK_REALTIME at boot）
    boot_time = current_realtime - uptime_seconds
    
    # 将 eBPF 时间戳转换为 Unix 时间戳（纳秒级）
    unix_ns = boot_time * 1e9 + bpf_ktime_ns
    return unix_ns

def print_event(cpu, data, size):
    event = b["events"].event(data)
    chrome_trace_event = {
        "name": f"Off-CPU Time for PID",
        "ph": "X",
        "pid": event.pid,
        "tid": event.pid,
        "ts": _convert_ktime_to_unix(event.start_time) / 1000.0,
        "dur":  (event.end_time - event.start_time) / 1000.0,
        "args": {}
    }	
    chrome_trace.append(chrome_trace_event)
	

if __name__ == "__main__":
    #import argparse
    #parser = argparse.ArgumentParser(description="Performance tracking script")
    #parser.add_argument("pid", type=int, help="Process ID to track")
    #args = parser.parse_args()
	
    filter_pid = sys.argv[1]

    # 替换TARGET_PID为实际的进程ID
    bpf_text = bpf_text.replace("TARGET_PID", str(filter_pid))
    b = BPF(text=bpf_text)

    # 注册tracepoint
    b.attach_kprobe(event_re="^finish_task_switch$|^finish_task_switch\.isra\.\d$", fn_name="trace_sched_switch")

    chrome_trace = []
    b["events"].open_perf_buffer(print_event)

    try:
        print(f"Tracking off-CPU time for PID . Press Ctrl+C to stop.")
        while True:
            b.kprobe_poll()
    except KeyboardInterrupt:
        pass

    # 将Chrome Trace格式的数据写入文件
    with open(f"offcpu_trace.json", "w") as f:
        json.dump(chrome_trace, f)

    print(f"Trace saved to offcpu_trace.json")



